import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.preprocessing import LabelEncoder
import pickle  # 使用 pickle 替代 joblib

class CrimePredictionModel:
    def __init__(self):
        # 初始化模型字典，包含随机森林和梯度提升树
        self.models = {
            'rf': RandomForestClassifier(n_estimators=50, max_depth=6, min_samples_leaf=5, random_state=42),
            'gb': GradientBoostingClassifier(n_estimators=50, max_depth=4, min_samples_leaf=5, random_state=42)
        }
        self.best_model = None  # 用于存储表现最好的模型
        self.best_model_name = None  # 最佳模型的名称
        self.label_encoder = LabelEncoder()  # 标签编码器，用于将犯罪类型转为数字
        
    def prepare_target(self, df):
        # 将犯罪类型（字符串）编码为数字标签，作为模型的目标变量
        return self.label_encoder.fit_transform(df['Crm Cd Desc'])
    
    def train_and_evaluate(self, X, y):
        # 划分训练集和测试集，比例为8:2
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42
        )
        results = {}
        # 遍历所有模型，分别训练和评估
        for name, model in self.models.items():
            model.fit(X_train, y_train)  # 训练模型
            train_score = model.score(X_train, y_train)  # 训练集准确率
            test_score = model.score(X_test, y_test)  # 测试集准确率
            cv_scores = cross_val_score(model, X, y, cv=5)  # 5折交叉验证得分
            y_pred = model.predict(X_test)  # 测试集预测结果
            results[name] = {
                'train_score': train_score,  # 训练集准确率
                'test_score': test_score,    # 测试集准确率
                'cv_mean': cv_scores.mean(), # 交叉验证平均分
                'cv_std': cv_scores.std(),   # 交叉验证标准差
                'classification_report': classification_report(y_test, y_pred),  # 分类报告
                'confusion_matrix': confusion_matrix(y_test, y_pred)  # 混淆矩阵
            }
        # 选择测试集准确率最高的模型作为最佳模型
        best_score = 0
        for name, result in results.items():
            if result['test_score'] > best_score:
                best_score = result['test_score']
                self.best_model = self.models[name]
                self.best_model_name = name
        return results
    
    def predict_hotspots(self, X, threshold=0.7):
        """
        预测犯罪热点
        """
        if self.best_model is None:
            raise ValueError("Model not trained yet!")
        
        # 获取预测概率
        probabilities = self.best_model.predict_proba(X)
        
        # 获取最高概率值
        max_probs = np.max(probabilities, axis=1)
        
        # 标识热点（概率高于阈值的位置）
        hotspots = max_probs > threshold
        
        return hotspots, max_probs
    
    def save_model(self, model_path):
        """
        保存模型
        """
        if self.best_model is None:
            raise ValueError("No model to save!")
        
        model_data = {
            'model': self.best_model,
            'model_name': self.best_model_name,
            'label_encoder': self.label_encoder
        }
        
        with open(model_path, 'wb') as f:
            pickle.dump(model_data, f)
        print(f"Model saved to {model_path}")
    
    @classmethod
    def load_model(cls, model_path):
        """
        加载保存的模型
        """
        with open(model_path, 'rb') as f:
            model_data = pickle.load(f)
        
        instance = cls()
        instance.best_model = model_data['model']
        instance.best_model_name = model_data['model_name']
        instance.label_encoder = model_data['label_encoder']
        
        return instance

def train_model(featured_path, processed_path, model_save_path):
    """
    训练模型的主函数
    """
    # 加载特征工程后的数据
    print("Loading data...")
    featured_data = pd.read_csv(featured_path)
    original_data = pd.read_csv(processed_path)
    
    # 创建模型实例
    model = CrimePredictionModel()
    
    # 准备目标变量
    y = model.prepare_target(original_data)
    
    # 训练和评估模型
    print("Training models...")
    results = model.train_and_evaluate(featured_data, y)
    
    # 保存最佳模型
    model.save_model(model_save_path)
    
    return model, results

if __name__ == "__main__":
    # 创建模型目录
    import os
    os.makedirs("data/models", exist_ok=True)
    
    # 训练模型
    model, results = train_model("data/processed/featured_crime_data.csv", "data/processed/processed_crime_data.csv", "data/models/best_crime_prediction_model.pkl")
    print("\nModel training completed successfully!") 